from diffusers import DiffusionPipeline
import torch
import os
import matplotlib.pyplot as plt

print(torch.cuda.is_available())
print(torch.__version__)
torch.cuda.empty_cache()

model_path = r'M:\moudels\stabilityai-sdxl-turbo'
cache_dir = r"F:\A_wokecode\gradio_study\pating\cachepath"

# 或者，更精细地控制模型各部分的位置
device_map = {"text_encoder": "cuda:0", "unet": "cuda:0", "vae": "cpu"}
pipe = DiffusionPipeline.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map=device_map,
    cache_dir=cache_dir,
    local_files_only=True

)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]

output_path = r"F:\A_wokecode\gradio_study\pating\outputs\generated_image.png"
os.makedirs(os.path.dirname(output_path), exist_ok=True)  # 确保目录存在
image.save(output_path)

# 使用 matplotlib 显示图片
plt.imshow(image)
plt.axis("off")  # 隐藏坐标轴
plt.show()
